import random,copy
import numpy as np
from tqdm import tqdm
import os

TRAINS_PATH = os.path.join(os.path.split(os.path.realpath(__file__))[0], '.\datas\\trains.txt')
FB15K_BASE_PATH = os.path.join(os.path.split(os.path.realpath(__file__))[0], '.\datas\\FB15K')
FB15K_BASE_PATH_1_1_PATH = os.path.join(FB15K_BASE_PATH, '1-1.txt')
FB15K_BASE_PATH_1_n_PATH = os.path.join(FB15K_BASE_PATH, '1-n.txt')
FB15K_BASE_PATH_n_1_PATH = os.path.join(FB15K_BASE_PATH, 'n-1.txt')
FB15K_BASE_PATH_n_n_PATH = os.path.join(FB15K_BASE_PATH, 'n-n.txt')
FB15K_BASE_PATH_test_PATH = os.path.join(FB15K_BASE_PATH, 'test2id.txt')
FB15K_BASE_PATH_train_PATH = os.path.join(FB15K_BASE_PATH, 'train2id.txt')
model_save_path=os.path.join(os.path.split(os.path.realpath(__file__))[0], '.\model_saved\\savedmodel.pkl')

def get_pairs(file_path1):
    relation=set()
    enity=set()
    pairs=[]
    with open(file_path1,'r') as f:
       for line in tqdm(f.readlines()):
           pair=line.strip().split()
           if len(pair)!=3:
               continue
           pairs.append(pair)
           enity.add(pair[0])
           enity.add(pair[1])
           relation.add(pair[2])
    enitylist=enity
    relationlist=relation
    return enitylist,relationlist,pairs

#测试数据加载与转换
def load_without_dic(pairs):
    return_pairs=[]
    
    for temp_pair in pairs:
        temp_data=[1]*3
        temp_data[0]=int(temp_pair[0])
        temp_data[1]=int(temp_pair[1])
        temp_data[2]=int(temp_pair[2])
        return_pairs.append(temp_data)
    return np.array(return_pairs)

class DataLoad():
    def __init__(self,enitylist,relationlist):
        self.enitylist=enitylist#实体list
        self.relationlist=relationlist#关系
        self.enity2id={v: i for i,v in enumerate(enitylist)}#生成index字典例如a[15]=10,生成index字典后b[10]=15
        self.relation2id={v: i for i,v in enumerate(relationlist)}#同上
      
        
      #当加载的数据不是已经转换为id的数据时，可以用下面的转换
    def get_index(self,dataset):
        return_dataset=[]
        
        for d in dataset:
            temp_data=[1]*3
            temp_data[0]=self.enity2id[d[0]]
            temp_data[1]=self.enity2id[d[1]]
            temp_data[2]=self.relation2id[d[2]]
            return_dataset.append(temp_data)
        return return_dataset
    
    #生成错误元组
    def __get_corrupt(self,dataset):
        corrupt_dataset=[]
        head=[]
        for triple in dataset:
            corrupt_pair=copy.deepcopy(triple)
            seed=random.random()
            if (seed>0.5):
                head=triple[0]
                while head==triple[0]:
                    head=random.sample(self.enitylist,1)[0]
                corrupt_pair[0]=head
            else:
                tail=triple[1]
                while tail==triple[1]:
                     tail=random.sample(self.enitylist,1)[0]
                corrupt_pair[1]=tail
            corrupt_dataset.append(corrupt_pair)
        return corrupt_dataset
    
    #数据迭代器-id数据
    def data_iter(self,pair,batchsize):
        for i in tqdm(range(len(pair)//batchsize)):
            correct_data=random.sample(pair,batchsize)
            corrupt_data=self.__get_corrupt(correct_data)
            #yield np.array(self.get_index(correct_data)),np.array(self.get_index(corrupt_data))
            yield np.array(correct_data,dtype=np.int32),np.array(corrupt_data,dtype=np.int32)
            
            
    #数据迭代器-非id数据    
    def data_iter_without_id(self,pair,batchsize):
        for i in tqdm(range(len(pair)//batchsize)):
            correct_data=random.sample(pair,batchsize)
            corrupt_data=self.__get_corrupt(correct_data)
            yield np.array(self.get_index(correct_data)),np.array(self.get_index(corrupt_data))
            #yield np.array(correct_data,dtype=np.int32),np.array(corrupt_data,dtype=np.int32)
                
    
    #非ID测试数据加载与转换
    def load_with_dic(self,test_pairs):
        return_data=[]
        for line in tqdm(test_pairs):
               temp_pair=[1]*3
               head=line[0]
               tail=line[1]
               relation=line[2]
               temp_pair[0]=self.enity2id[head]
               temp_pair[1]=self.enity2id[tail]
               temp_pair[2]=self.relation2id[relation]
               return_data.append((temp_pair))
        return np.array(return_data)


				
if __name__=='__main__':
    #genrateParis()
    enity, relationShips, pairs=get_pairs(FB15K_BASE_PATH_train_PATH)
    dataload=DataLoad(enity, relationShips)
    test_data=dataload.load_with_dic(FB15K_BASE_PATH_test_PATH)
    i=1